/*
 * Copyright (c) 2022-2024 KCloud-Platform-IoT Author or Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.rocketmq.client.impl.consumer;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.*;

import lombok.Getter;
import org.apache.rocketmq.client.consumer.DefaultMQPushConsumer;
import org.apache.rocketmq.client.consumer.listener.ConsumeConcurrentlyContext;
import org.apache.rocketmq.client.consumer.listener.ConsumeConcurrentlyStatus;
import org.apache.rocketmq.client.consumer.listener.ConsumeReturnType;
import org.apache.rocketmq.client.consumer.listener.MessageListenerConcurrently;
import org.apache.rocketmq.client.hook.ConsumeMessageContext;
import org.apache.rocketmq.client.stat.ConsumerStatsManager;
import org.apache.rocketmq.common.MixAll;
import org.apache.rocketmq.common.ThreadFactoryImpl;
import org.apache.rocketmq.common.UtilAll;
import org.apache.rocketmq.common.message.MessageAccessor;
import org.apache.rocketmq.common.message.MessageExt;
import org.apache.rocketmq.common.message.MessageQueue;
import org.apache.rocketmq.common.utils.ThreadUtils;
import org.apache.rocketmq.remoting.protocol.body.CMResult;
import org.apache.rocketmq.remoting.protocol.body.ConsumeMessageDirectlyResult;
import org.apache.rocketmq.logging.org.slf4j.Logger;
import org.apache.rocketmq.logging.org.slf4j.LoggerFactory;
import org.laokou.common.core.config.VirtualThreadFactory;

/**
 * @author laokou
 */
public class ConsumeMessageConcurrentlyService implements ConsumeMessageService {

	private static final Logger log = LoggerFactory.getLogger(ConsumeMessageConcurrentlyService.class);

	private final DefaultMQPushConsumerImpl defaultMQPushConsumerImpl;

	private final DefaultMQPushConsumer defaultMQPushConsumer;

	private final MessageListenerConcurrently messageListener;

	private final ThreadPoolExecutor consumeExecutor;

	private final String consumerGroup;

	private final ScheduledExecutorService scheduledExecutorService;

	private final ScheduledExecutorService cleanExpireMsgExecutors;

	public ConsumeMessageConcurrentlyService(DefaultMQPushConsumerImpl defaultMQPushConsumerImpl,
			MessageListenerConcurrently messageListener) {
		this.defaultMQPushConsumerImpl = defaultMQPushConsumerImpl;
		this.messageListener = messageListener;

		this.defaultMQPushConsumer = this.defaultMQPushConsumerImpl.getDefaultMQPushConsumer();
		this.consumerGroup = this.defaultMQPushConsumer.getConsumerGroup();
		BlockingQueue<Runnable> consumeRequestQueue = new LinkedBlockingQueue<>();

		String consumerGroupTag = (consumerGroup.length() > 100 ? consumerGroup.substring(0, 100) : consumerGroup)
				+ "_";

		// 虚拟线程
		this.consumeExecutor = new ThreadPoolExecutor(0, // corePoolSize为0，因为不需要保留核心线程
				Integer.MAX_VALUE, // maximumPoolSize: 无限制，允许任意数量的虚拟线程
				60L, // keepAliveTime: 虚拟线程不需要保持存活时间
				TimeUnit.SECONDS, // 时间单位
				new SynchronousQueue<>(), // 使用SynchronousQueue以确保每个任务都会创建一个新线程
				VirtualThreadFactory.INSTANCE // 使用虚拟线程工厂
		);

		this.scheduledExecutorService = Executors.newSingleThreadScheduledExecutor(
				new ThreadFactoryImpl("ConsumeMessageScheduledThread_" + consumerGroupTag));
		this.cleanExpireMsgExecutors = Executors.newSingleThreadScheduledExecutor(
				new ThreadFactoryImpl("CleanExpireMsgScheduledThread_" + consumerGroupTag));
	}

	public void start() {
		this.cleanExpireMsgExecutors.scheduleAtFixedRate(new Runnable() {

			@Override
			public void run() {
				try {
					cleanExpireMsg();
				}
				catch (Throwable e) {
					log.error("scheduleAtFixedRate cleanExpireMsg exception", e);
				}
			}

		}, this.defaultMQPushConsumer.getConsumeTimeout(), this.defaultMQPushConsumer.getConsumeTimeout(),
				TimeUnit.MINUTES);
	}

	public void shutdown(long awaitTerminateMillis) {
		this.scheduledExecutorService.shutdown();
		ThreadUtils.shutdownGracefully(this.consumeExecutor, awaitTerminateMillis, TimeUnit.MILLISECONDS);
		this.cleanExpireMsgExecutors.shutdown();
	}

	@Override
	public void updateCorePoolSize(int corePoolSize) {
		if (corePoolSize > 0 && corePoolSize <= Short.MAX_VALUE
				&& corePoolSize < this.defaultMQPushConsumer.getConsumeThreadMax()) {
			this.consumeExecutor.setCorePoolSize(corePoolSize);
		}
	}

	@Override
	public void incCorePoolSize() {

	}

	@Override
	public void decCorePoolSize() {

	}

	@Override
	public int getCorePoolSize() {
		return this.consumeExecutor.getCorePoolSize();
	}

	@Override
	public ConsumeMessageDirectlyResult consumeMessageDirectly(MessageExt msg, String brokerName) {
		ConsumeMessageDirectlyResult result = new ConsumeMessageDirectlyResult();
		result.setOrder(false);
		result.setAutoCommit(true);

		msg.setBrokerName(brokerName);
		List<MessageExt> msgs = new ArrayList<>();
		msgs.add(msg);
		MessageQueue mq = new MessageQueue();
		mq.setBrokerName(brokerName);
		mq.setTopic(msg.getTopic());
		mq.setQueueId(msg.getQueueId());

		ConsumeConcurrentlyContext context = new ConsumeConcurrentlyContext(mq);

		this.defaultMQPushConsumerImpl.resetRetryAndNamespace(msgs, this.consumerGroup);

		final long beginTime = System.currentTimeMillis();

		log.info("consumeMessageDirectly receive new message: {}", msg);

		try {
			ConsumeConcurrentlyStatus status = this.messageListener.consumeMessage(msgs, context);
			if (status != null) {
				switch (status) {
					case CONSUME_SUCCESS:
						result.setConsumeResult(CMResult.CR_SUCCESS);
						break;
					case RECONSUME_LATER:
						result.setConsumeResult(CMResult.CR_LATER);
						break;
					default:
						break;
				}
			}
			else {
				result.setConsumeResult(CMResult.CR_RETURN_NULL);
			}
		}
		catch (Throwable e) {
			result.setConsumeResult(CMResult.CR_THROW_EXCEPTION);
			result.setRemark(UtilAll.exceptionSimpleDesc(e));

			log.warn("consumeMessageDirectly exception: {} Group: {} Msgs: {} MQ: {}", UtilAll.exceptionSimpleDesc(e),
					ConsumeMessageConcurrentlyService.this.consumerGroup, msgs, mq, e);
		}

		result.setSpentTimeMills(System.currentTimeMillis() - beginTime);

		log.info("consumeMessageDirectly Result: {}", result);

		return result;
	}

	@Override
	public void submitConsumeRequest(final List<MessageExt> msgs, final ProcessQueue processQueue,
			final MessageQueue messageQueue, final boolean dispatchToConsume) {
		final int consumeBatchSize = this.defaultMQPushConsumer.getConsumeMessageBatchMaxSize();
		if (msgs.size() <= consumeBatchSize) {
			ConsumeRequest consumeRequest = new ConsumeRequest(msgs, processQueue, messageQueue);
			try {
				this.consumeExecutor.submit(consumeRequest);
			}
			catch (RejectedExecutionException e) {
				this.submitConsumeRequestLater(consumeRequest);
			}
		}
		else {
			for (int total = 0; total < msgs.size();) {
				List<MessageExt> msgThis = new ArrayList<>(consumeBatchSize);
				for (int i = 0; i < consumeBatchSize; i++, total++) {
					if (total < msgs.size()) {
						msgThis.add(msgs.get(total));
					}
					else {
						break;
					}
				}

				ConsumeRequest consumeRequest = new ConsumeRequest(msgThis, processQueue, messageQueue);
				try {
					this.consumeExecutor.submit(consumeRequest);
				}
				catch (RejectedExecutionException e) {
					for (; total < msgs.size(); total++) {
						msgThis.add(msgs.get(total));
					}

					this.submitConsumeRequestLater(consumeRequest);
				}
			}
		}
	}

	@Override
	public void submitPopConsumeRequest(final List<MessageExt> msgs, final PopProcessQueue processQueue,
			final MessageQueue messageQueue) {
		throw new UnsupportedOperationException();
	}

	private void cleanExpireMsg() {
		Iterator<Map.Entry<MessageQueue, ProcessQueue>> it = this.defaultMQPushConsumerImpl.getRebalanceImpl()
			.getProcessQueueTable()
			.entrySet()
			.iterator();
		while (it.hasNext()) {
			Map.Entry<MessageQueue, ProcessQueue> next = it.next();
			ProcessQueue pq = next.getValue();
			pq.cleanExpiredMsg(this.defaultMQPushConsumer);
		}
	}

	public void processConsumeResult(final ConsumeConcurrentlyStatus status, final ConsumeConcurrentlyContext context,
			final ConsumeRequest consumeRequest) {
		int ackIndex = context.getAckIndex();

		if (consumeRequest.getMsgs().isEmpty())
			return;

		switch (status) {
			case CONSUME_SUCCESS:
				if (ackIndex >= consumeRequest.getMsgs().size()) {
					ackIndex = consumeRequest.getMsgs().size() - 1;
				}
				int ok = ackIndex + 1;
				int failed = consumeRequest.getMsgs().size() - ok;
				this.getConsumerStatsManager()
					.incConsumeOKTPS(consumerGroup, consumeRequest.getMessageQueue().getTopic(), ok);
				this.getConsumerStatsManager()
					.incConsumeFailedTPS(consumerGroup, consumeRequest.getMessageQueue().getTopic(), failed);
				break;
			case RECONSUME_LATER:
				ackIndex = -1;
				this.getConsumerStatsManager()
					.incConsumeFailedTPS(consumerGroup, consumeRequest.getMessageQueue().getTopic(),
							consumeRequest.getMsgs().size());
				break;
			default:
				break;
		}

		switch (this.defaultMQPushConsumer.getMessageModel()) {
			case BROADCASTING:
				for (int i = ackIndex + 1; i < consumeRequest.getMsgs().size(); i++) {
					MessageExt msg = consumeRequest.getMsgs().get(i);
					log.warn("BROADCASTING, the message consume failed, drop it, {}", msg.toString());
				}
				break;
			case CLUSTERING:
				List<MessageExt> msgBackFailed = new ArrayList<>(consumeRequest.getMsgs().size());
				for (int i = ackIndex + 1; i < consumeRequest.getMsgs().size(); i++) {
					MessageExt msg = consumeRequest.getMsgs().get(i);
					// Maybe message is expired and cleaned, just ignore it.
					if (!consumeRequest.getProcessQueue().containsMessage(msg)) {
						log.info(
								"Message is not found in its process queue; skip send-back-procedure, topic={}, "
										+ "brokerName={}, queueId={}, queueOffset={}",
								msg.getTopic(), msg.getBrokerName(), msg.getQueueId(), msg.getQueueOffset());
						continue;
					}
					boolean result = this.sendMessageBack(msg, context);
					if (!result) {
						msg.setReconsumeTimes(msg.getReconsumeTimes() + 1);
						msgBackFailed.add(msg);
					}
				}

				if (!msgBackFailed.isEmpty()) {
					consumeRequest.getMsgs().removeAll(msgBackFailed);

					this.submitConsumeRequestLater(msgBackFailed, consumeRequest.getProcessQueue(),
							consumeRequest.getMessageQueue());
				}
				break;
			default:
				break;
		}

		long offset = consumeRequest.getProcessQueue().removeMessage(consumeRequest.getMsgs());
		if (offset >= 0 && !consumeRequest.getProcessQueue().isDropped()) {
			this.defaultMQPushConsumerImpl.getOffsetStore()
				.updateOffset(consumeRequest.getMessageQueue(), offset, true);
		}
	}

	public ConsumerStatsManager getConsumerStatsManager() {
		return this.defaultMQPushConsumerImpl.getConsumerStatsManager();
	}

	public boolean sendMessageBack(final MessageExt msg, final ConsumeConcurrentlyContext context) {
		int delayLevel = context.getDelayLevelWhenNextConsume();

		// Wrap topic with namespace before sending back message.
		msg.setTopic(this.defaultMQPushConsumer.withNamespace(msg.getTopic()));
		try {
			this.defaultMQPushConsumerImpl.sendMessageBack(msg, delayLevel,
					this.defaultMQPushConsumer.queueWithNamespace(context.getMessageQueue()));
			return true;
		}
		catch (Exception e) {
			log.error("sendMessageBack exception, group: " + this.consumerGroup + " msg: " + msg, e);
		}

		return false;
	}

	private void submitConsumeRequestLater(final List<MessageExt> msgs, final ProcessQueue processQueue,
			final MessageQueue messageQueue) {

		this.scheduledExecutorService.schedule(() -> ConsumeMessageConcurrentlyService.this.submitConsumeRequest(msgs,
				processQueue, messageQueue, true), 5000, TimeUnit.MILLISECONDS);
	}

	private void submitConsumeRequestLater(final ConsumeRequest consumeRequest) {

		this.scheduledExecutorService.schedule(() -> {
			ConsumeMessageConcurrentlyService.this.consumeExecutor.submit(consumeRequest);
		}, 5000, TimeUnit.MILLISECONDS);
	}

	@Getter
	class ConsumeRequest implements Runnable {

		private final List<MessageExt> msgs;

		private final ProcessQueue processQueue;

		private final MessageQueue messageQueue;

		public ConsumeRequest(List<MessageExt> msgs, ProcessQueue processQueue, MessageQueue messageQueue) {
			this.msgs = msgs;
			this.processQueue = processQueue;
			this.messageQueue = messageQueue;
		}

		@Override
		public void run() {
			if (this.processQueue.isDropped()) {
				log.info("the message queue not be able to consume, because it's dropped. group={} {}",
						ConsumeMessageConcurrentlyService.this.consumerGroup, this.messageQueue);
				return;
			}

			ConsumeConcurrentlyContext context = new ConsumeConcurrentlyContext(messageQueue);
			ConsumeConcurrentlyStatus status = null;
			defaultMQPushConsumerImpl.tryResetPopRetryTopic(msgs, consumerGroup);
			defaultMQPushConsumerImpl.resetRetryAndNamespace(msgs, defaultMQPushConsumer.getConsumerGroup());

			ConsumeMessageContext consumeMessageContext = null;
			if (ConsumeMessageConcurrentlyService.this.defaultMQPushConsumerImpl.hasHook()) {
				consumeMessageContext = new ConsumeMessageContext();
				consumeMessageContext.setNamespace(defaultMQPushConsumer.getNamespace());
				consumeMessageContext.setConsumerGroup(defaultMQPushConsumer.getConsumerGroup());
				consumeMessageContext.setProps(new HashMap<>());
				consumeMessageContext.setMq(messageQueue);
				consumeMessageContext.setMsgList(msgs);
				consumeMessageContext.setSuccess(false);
				ConsumeMessageConcurrentlyService.this.defaultMQPushConsumerImpl
					.executeHookBefore(consumeMessageContext);
			}

			long beginTimestamp = System.currentTimeMillis();
			boolean hasException = false;
			ConsumeReturnType returnType = ConsumeReturnType.SUCCESS;
			try {
				if (msgs != null && !msgs.isEmpty()) {
					for (MessageExt msg : msgs) {
						MessageAccessor.setConsumeStartTimeStamp(msg, String.valueOf(System.currentTimeMillis()));
					}
				}
				status = ConsumeMessageConcurrentlyService.this.messageListener
					.consumeMessage(Collections.unmodifiableList(msgs), context);
			}
			catch (Throwable e) {
				log.warn("consumeMessage exception: {} Group: {} Msgs: {} MQ: {}", UtilAll.exceptionSimpleDesc(e),
						ConsumeMessageConcurrentlyService.this.consumerGroup, msgs, messageQueue, e);
				hasException = true;
			}
			long consumeRT = System.currentTimeMillis() - beginTimestamp;
			if (null == status) {
				if (hasException) {
					returnType = ConsumeReturnType.EXCEPTION;
				}
				else {
					returnType = ConsumeReturnType.RETURNNULL;
				}
			}
			else if (consumeRT >= defaultMQPushConsumer.getConsumeTimeout() * 60 * 1000) {
				returnType = ConsumeReturnType.TIME_OUT;
			}
			else if (ConsumeConcurrentlyStatus.RECONSUME_LATER == status) {
				returnType = ConsumeReturnType.FAILED;
			}
			else if (ConsumeConcurrentlyStatus.CONSUME_SUCCESS == status) {
				returnType = ConsumeReturnType.SUCCESS;
			}

			if (ConsumeMessageConcurrentlyService.this.defaultMQPushConsumerImpl.hasHook()) {
				consumeMessageContext.getProps().put(MixAll.CONSUME_CONTEXT_TYPE, returnType.name());
			}

			if (null == status) {
				log.warn("consumeMessage return null, Group: {} Msgs: {} MQ: {}",
						ConsumeMessageConcurrentlyService.this.consumerGroup, msgs, messageQueue);
				status = ConsumeConcurrentlyStatus.RECONSUME_LATER;
			}

			if (ConsumeMessageConcurrentlyService.this.defaultMQPushConsumerImpl.hasHook()) {
				consumeMessageContext.setStatus(status.toString());
				consumeMessageContext.setSuccess(ConsumeConcurrentlyStatus.CONSUME_SUCCESS == status);
				consumeMessageContext.setAccessChannel(defaultMQPushConsumer.getAccessChannel());
				ConsumeMessageConcurrentlyService.this.defaultMQPushConsumerImpl
					.executeHookAfter(consumeMessageContext);
			}

			ConsumeMessageConcurrentlyService.this.getConsumerStatsManager()
				.incConsumeRT(ConsumeMessageConcurrentlyService.this.consumerGroup, messageQueue.getTopic(), consumeRT);

			if (!processQueue.isDropped()) {
				ConsumeMessageConcurrentlyService.this.processConsumeResult(status, context, this);
			}
			else {
				log.warn("processQueue is dropped without process consume result. messageQueue={}, msgs={}",
						messageQueue, msgs);
			}
		}

	}

}
